import logging
import torch
import numpy as np
from haversine import haversine

def extract_user_loc(prob_mat):
    prob_mat1 = prob_mat.cpu()
    pred_loc = torch.max(prob_mat1, 1)[1]
    pred_loc_list = pred_loc.numpy().tolist()
    return pred_loc_list

def evaluate(y_pred, y_true, classLatMedian, classLonMedian):
    assert len(y_pred) == len(y_true), "#preds: %d, #users: %d" % (len(y_pred), len(y_true))
    y_pred = extract_user_loc(y_pred)
    distances = []
    latlon_pred = []
    latlon_true = []

    for i in range(len(y_pred)):
        # 真实位置
        location = y_true[i]
        lat, lon = float(classLatMedian[str(location)]), float(classLonMedian[str(location)])
        latlon_true.append([lat, lon])

        prediction = y_pred[i]
        if prediction == -1:
            continue

        # 预测位置
        lat_pred, lon_pred = float(classLatMedian[str(prediction)]), float(classLonMedian[str(prediction)])
        latlon_pred.append([lat_pred, lon_pred])
        distance = haversine((lat, lon), (lat_pred, lon_pred))
        distances.append(distance)

    acc_at_161 = float(len([d for d in distances if d < 161]) / float(len(distances)))
    logging.info(
        "Mean: " + str(int(np.mean(distances))) + " Median: " + str(int(np.median(distances))) + " Acc@161: " + str(
            int(acc_at_161)))
    # print("Mean: " + str(int(np.mean(distances))) + " Median: " + str(int(np.median(distances))) + " Acc@161: " + str(
    #         (acc_at_161)))

    # return np.mean(distances), np.median(distances), acc_at_161, distances, latlon_true, latlon_pred
    return acc_at_161, np.median(distances), np.mean(distances)


def cal_acc(predictions, ground_truth):
    if len(predictions) == len(ground_truth):
        correct = 0
        eval_num = len(predictions)
        for i in range(eval_num):
            if predictions[i] == ground_truth[i]:
                correct += 1
        acc = float(correct / eval_num)
        print('acc:%f' % acc)
    else:
        return -1